/**
 * 版权所有 (c) 2018，《自制AI图像搜索引擎》作者。

 * 您可以下载和使用该软件。但它仅供个人学习和研究使用，禁止任何商业用途。

 * 您必须在此软件的所有副本（包括任何修改或衍生版本）中保留此许可证文件，且不能对其有任何修改。

 * 欢迎购买《自制AI图像搜索引擎》——一本全面介绍AI图像搜索引擎原理与实现的书。如有任何意见或建议，请发送邮件至imgsearch@126.com。
 */
package com.ai.deepsearch;

import org.apache.commons.cli.*;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.util.ModelSerializer;
import org.mapdb.DB;
import org.mapdb.DBMaker;
import org.mapdb.Serializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.VGG16ImagePreProcessor;

import java.io.File;
import java.io.IOException;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.ConcurrentMap;

/**
 * 图像特征库生成工具
 */
public class GenerateImgsFeatDBTool {
    private DB db;
    private ConcurrentMap<String, double[]> map;
    private static ComputationGraph vgg16Model;

    private void initDB(String dbName) {
        System.out.println("GenerateImgsFeatDBTool init db");
        db = DBMaker.fileDB(dbName).fileMmapEnable().make();
        map = db.hashMap("feat_map", Serializer.STRING, Serializer.DOUBLE_ARRAY).createOrOpen();
    }

    private void loadVGGModel(String modelFilePath) throws IOException {
        File vgg16ModelFile = new File(modelFilePath);
        vgg16Model = ModelSerializer.restoreComputationGraph(vgg16ModelFile);
    }

    private INDArray getImgFeature(File imgFile) throws IOException {
        NativeImageLoader loader = new NativeImageLoader(224, 224, 3);
        INDArray imageArray = loader.asMatrix(imgFile);
        DataNormalization scaler = new VGG16ImagePreProcessor();
        scaler.transform(imageArray);
        Map<String, INDArray> map = vgg16Model.feedForward(imageArray, false);
        INDArray feature = map.get("fc2");
        return feature;
    }

    private double[] INDArray2DoubleArray(INDArray indArr) {
        String indArrStr = indArr.toString().replace("[", "").replace("]", "");
        String[] strArr = indArrStr.split(",");
        int len = strArr.length;
        double[] doubleArr = new double[len];
        for (int i = 0; i < len; i++) {
            doubleArr[i] = Double.parseDouble(strArr[i]);
        }
        return doubleArr;
    }

    private void exportImgsFeature2DB(String imgDirName) throws IOException {
        File dir = new File(imgDirName);
        if (dir.isDirectory()) {
            File[] fileList = dir.listFiles();
            int len = fileList.length;
            for (int i = 0; i < len; i++) {
                INDArray feat = getImgFeature(fileList[i]);
                double[] featD = INDArray2DoubleArray(feat);
                String imgName = fileList[i].getName();
                System.out.println("GenerateImgsFeatDBTool:" + i + ":" + imgName + "," + featD.toString());
                map.put(imgName, featD);
            }
        } else {
            System.out.println(imgDirName + "is not a directory!");
        }
        db.close();
        System.out.println("GenerateImgsFeatDBTool close db");
    }

    public static void main(String[] args) throws IOException {
        String usage = "java -jar GenerateImgsFeatDBTool.jar [-h] -m 模型路径名 -d 特征库路径名 -i 图像文件夹路径名";
        HelpFormatter formatter = new HelpFormatter();
        formatter.setWidth(200);
        CommandLineParser parser = new DefaultParser();

        Option help = new Option("h",false,"显示帮助信息");
        Option model=Option.builder("m").hasArg().argName("model").desc("模型路径名").build();
        Option database = Option.builder("d").hasArg().argName("database").desc("图像特征库路径名").build();
        Option img = Option.builder("i").hasArg().argName("imgdir").desc("用于构建特征库的图像文件夹路径全名").build();
        Option dist = Option.builder("t").hasArg().argName("dist").desc("分类后的目标文件夹").build();
        Option useDb = Option.builder("b").hasArg().argName("usedb").desc("使用现成的db").build();

        Options options = new Options();
        options.addOption(help);
        options.addOption(model);
        options.addOption(database);
        options.addOption(img);
        options.addOption(dist);
        options.addOption(useDb);

        String modelFilePath=null;
        String dbFilePath = null;
        String imgsDir = null;
        String distDir=null;
        boolean isUseDb=false;
        Integer k=null;

        try {
            CommandLine line = parser.parse(options, args);
            if (line.getOptions().length > 0) {
                if (line.hasOption("h")) {
                    formatter.printHelp(usage, options);
                }
                if (line.hasOption("m")) {
                    modelFilePath = line.getOptionValue("m");
                }
                if((!line.hasOption("h"))&&(!line.hasOption("m"))) {
                    System.out.println("缺少模型参数!");
                }
                if (line.hasOption("d")) {
                    dbFilePath = line.getOptionValue("d");
                }
                if(line.hasOption("t")){
                    distDir=line.getOptionValue("t");
                }

                if(line.hasOption("b")){
                    isUseDb=true;
                    String kStr = line.getOptionValue("b");
                    k=Integer.parseInt(kStr);
                }

                if((!line.hasOption("h"))&&(!line.hasOption("d"))) {
                    System.out.println("缺少特征库参数!");
                }
                if (line.hasOption("i")) {
                    imgsDir = line.getOptionValue("i");
                }
                if((!line.hasOption("h"))&&(!line.hasOption("i"))) {
                    System.out.println("缺少图像文件夹参数!");
                }
            } else {
                System.out.println("参数为空!");
            }
        } catch (ParseException e) {
            String message=e.getMessage();
            String[] messages=message.split(":");
            if("Missing argument for option".equals(messages[0])) {
                if("m".equals(messages[1].trim())) {
                    System.out.println("缺少模型参数值!");
                }
                if("d".equals(messages[1].trim())) {
                    System.out.println("缺少特征库参数值!");
                }
                if("".equals(messages[1].trim())) {
                    System.out.println("缺少图像文件夹参数值!");
                }
            }
        }

        System.out.println("开始生成向量 "+new Date());

        if (!isUseDb&&(modelFilePath != null)&&(dbFilePath != null) && (imgsDir != null)) {
            GenerateImgsFeatDBTool tool = new GenerateImgsFeatDBTool();
            tool.initDB(dbFilePath);
            try {
                tool.loadVGGModel(modelFilePath);
                tool.exportImgsFeature2DB(imgsDir);
            } catch (IOException e) {
                e.printStackTrace();
                System.out.println("GenerateImgsFeatDBTool遇到错误" + e.getMessage());
            }
        }else {
            System.out.println("跳过生成db 使用文件 "+dbFilePath);
        }

        System.out.println("向量生成完毕 "+new Date());

        if(distDir!=null){
            Classify.kmeans(dbFilePath,imgsDir,distDir,k);
        }

    }
}
